package com.moxc.learn_springai_01.utils;

import org.springframework.ai.chat.metadata.DefaultUsage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatResponse;

public final class UsageUtils {
    private UsageUtils() {
        throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
    }

    public static Usage getCumulativeUsage(final Usage currentUsage, final ChatResponse previousChatResponse) {
        Usage usageFromPreviousChatResponse = null;
        if (previousChatResponse != null && previousChatResponse.getMetadata() != null && previousChatResponse.getMetadata().getUsage() != null) {
            usageFromPreviousChatResponse = previousChatResponse.getMetadata().getUsage();
            if (!isEmpty(currentUsage)) {
                Integer promptTokens = currentUsage.getPromptTokens();
                Integer generationTokens = currentUsage.getCompletionTokens();
                Integer totalTokens = currentUsage.getTotalTokens();
                promptTokens = promptTokens + usageFromPreviousChatResponse.getPromptTokens();
                generationTokens = generationTokens + usageFromPreviousChatResponse.getCompletionTokens();
                totalTokens = totalTokens + usageFromPreviousChatResponse.getTotalTokens();
                return new DefaultUsage(promptTokens, generationTokens, totalTokens);
            } else {
                return usageFromPreviousChatResponse;
            }
        } else {
            return currentUsage;
        }
    }

    public static boolean isEmpty(Usage usage) {
        if (usage == null) {
            return true;
        } else {
            return usage != null && (long)usage.getTotalTokens() == 0L;
        }
    }
}